import asyncio
from typing import TypedDict

from langchain_community.chat_models import ChatTongyi
from langgraph.graph import START, StateGraph

import os
os.environ["DASHSCOPE_API_KEY"] = os.getenv("OPENAI_API_KEY")

# 打上标签类型
joke_model = ChatTongyi(model="qwen-plus", tags=["joke"])
poem_model = ChatTongyi(model="qwen-plus", tags=["poem"])


class State(TypedDict):
      topic: str
      joke: str
      poem: str


async def call_model(state, config):
      topic = state["topic"]
      print("Writing joke...")

      joke_response = await joke_model.ainvoke(
            [{"role": "user", "content": f"Write a joke about {topic}"}],
            config,
      )
      print("\n\nWriting poem...")
      poem_response = await poem_model.ainvoke(
            [{"role": "user", "content": f"Write a short poem about {topic}"}],
            config,
      )
      return {"joke": joke_response.content, "poem": poem_response.content}


graph = (
      StateGraph(State)
      .add_node(call_model)
      .add_edge(START, "call_model")
      .compile()
)

async def main():
    results = []
    async for msg, metadata in graph.astream(
            {"topic": "cats"},
            stream_mode="messages",
    ):
        if metadata["tags"] == ["joke"]:
            print(msg.content, end="|", flush=True)
            # 收集数据
            results.append((msg, metadata))
    return results

if __name__ == '__main__':
    result = asyncio.run(main())
    print("\nReceived", len(result), "messages")

